# !/usr/bin/python

from Crypto.Cipher import AES
import base64


class Aescrypt:
    def __init__(self, key: str, model: str, iv: str, encode_='utf-8', pad='ZeroPadding'):
        self.encode_ = encode_
        self.model = {'ECB': AES.MODE_ECB, 'CBC': AES.MODE_CBC}[model]
        self.pad = pad
        self.padding_dict = {
            "ZeroPadding": [self.zero_padding, self.clean_result],
            "PKCS5Padding": [self.pkcs5_padding, self.pkcs5_unpadding],
        }
        # 密钥必须补位到16位
        self.key = self.zero_padding(key)

        if model == 'ECB':
            self.aes = AES.new(self.key, self.model)  # 创建一个aes对象
        elif model == 'CBC':
            self.aes = AES.new(self.key, self.model, iv.encode())  # 创建一个aes对象

        # 这里的密钥长度必须是16、24或32，目前16位的就够用了

    def zero_padding(self, par):
        par = par.encode(self.encode_)
        while len(par) % AES.block_size != 0:
            par += b"\0"
        return par

    @staticmethod
    def pkcs5_padding(data):
        pad = lambda s: s + (AES.block_size - len(s.encode()) % AES.block_size) * chr(AES.block_size - len(s.encode()) % AES.block_size)

        return pad(data).encode()

    @staticmethod
    def pkcs5_unpadding(padded_data):
        unpad = lambda s: s[:-ord(s[len(s) - 1:])]
        return unpad(padded_data)

    @staticmethod
    def clean_result(value):
        res_list = []
        for c in value:
            # ascii码范围获取
            if ord(c) > 31 & ord(c) < 127:
                res_list.append(c)
        return ''.join(res_list)

    def aesencrypt(self, text):
        tmp = self.padding_dict[self.pad][0](text)
        self.encrypt_text = self.aes.encrypt(tmp)
        return base64.encodebytes(self.encrypt_text).decode().strip()

    def aesdecrypt(self, text):
        if not text:
            return text
        tmp = base64.decodebytes(text.encode(self.encode_))
        self.decrypt_text = self.aes.decrypt(tmp)
        tmp = self.padding_dict[self.pad][1](self.decrypt_text.decode(self.encode_))
        return tmp


def db_aesdecrypt(db_key, t):
    """
    db解密函数
    :param t:
    :return:
    """
    if not db_key or not t:
        return t
    return Aescrypt(db_key, "CBC", db_key, "utf-8", pad="PKCS5Padding").aesdecrypt(t)


if __name__ == '__main__':

    # aes_cbc = Aescrypt("1234567890123456", "CBC", "1234567890123456", "utf-8", pad="PKCS5Padding")
    # test = 'qianyu'
    # tmp = aes_cbc.aesencrypt(test)
    # print(tmp)
    #
    # res = Aescrypt("1234567890123456", "CBC", "1234567890123456", "utf-8", pad="PKCS5Padding").aesdecrypt(tmp)
    # print(test == res)

    AES_KEY_ecb = 'XXLhYQGwHEvZIzd'
    aes_ecb = Aescrypt(AES_KEY_ecb, "ECB", "", "utf-8")
    data = "qianyu"
    tmp = aes_ecb.aesencrypt(data)
    # 密⽂
    print("密文", tmp)
    #
    res = aes_ecb.aesdecrypt(tmp)
    print('明⽂:', res)